package plat.wx.context;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import plat.wx.config.MyConfig;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 上下文保存
 *
 * @author wangyy
 */
@Component
public class Session {

    private static Map<String, List<Map<String, String>>> userSession = new HashMap<>();

    @Autowired
    MyConfig config;

    /**
     * 构建包含上下文的请求内容
     *
     * @param query
     * @param userId
     * @return
     */
    public String buildSessionQuery(String query, String userId) {
        String prompt = config.getCharacterDesc();

        if (StringUtils.hasLength(prompt))
            prompt += "\n\n";

        List<Map<String, String>> session = userSession.get(userId);
        if (session != null) {
            for (int i = 0; i < session.size(); i++) {
                Map<String, String> map = session.get(i);
                prompt += "Q: " + map.get("question") + "\n\nA: " + map.get("answer") + "\n\n";
            }
            return prompt + ("Q: " + query + "\n\nA: ");
        } else {
            return prompt + "Q: " + query + "\n\nA: ";
        }
    }

    /**
     * 保存聊天记录
     *
     * @param query
     * @param answer
     * @param userId
     */
    public void saveSession(String query, String answer, String userId) {

        Integer conversationMaxTokens = config.getConversationMaxTokens();

        Map<String, String> conversation = new HashMap<>();
        conversation.put("question", query);
        conversation.put("answer", answer);

        List<Map<String, String>> maps = userSession.get(userId);

        if (maps != null) {
            maps.add(conversation);
        } else {
            List<Map<String, String>> queue = new ArrayList<>();
            queue.add(conversation);
            userSession.put(userId, queue);
        }

        discardDxceedConversation(userSession.get(userId), conversationMaxTokens);
    }

    /**
     * 检查聊天记录上下文是否超限，如果超限，删除最早一次谈话记录
     *
     * @param session
     * @param maxToken
     */
    private void discardDxceedConversation(List<Map<String, String>> session, Integer maxToken) {

        AtomicInteger count = new AtomicInteger();
        session.forEach(item -> {
            String question = item.get("question");
            String answer = item.get("answer");
            count.getAndAdd(question.length() + answer.length());
        });

        if (count.get() > maxToken)
            session.remove(0);
    }

    public void clearSession(String userId) {
        if (StringUtils.hasLength(userId))
            userSession.remove(userId);
    }
}